package com.chukun.flink.dataset.operator;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * @author chukun
 * @version 1.0.0
 * @description coGroup操作符将两个数据集合并为一个数据集
 * @createTime 2022年05月28日 23:39:00
 */
public class CoGroupOperator {

    public static void main(String[] args) throws Exception {

        List<Tuple2<Integer, String>> commodityList = new ArrayList<>();
        List<Tuple2<String, Integer>> orderList = new ArrayList<>();

        commodityList.add(new Tuple2<>(1, "手机"));
        commodityList.add(new Tuple2<>(2, "电脑"));
        commodityList.add(new Tuple2<>(3, "移动电源"));

        orderList.add(new Tuple2<>("张三", 1));
        orderList.add(new Tuple2<>("赵六", 1));
        orderList.add(new Tuple2<>("李四", 2));
        orderList.add(new Tuple2<>("刘七", 3));
        orderList.add(new Tuple2<>("王五", 2));

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSet<Tuple2<Integer, String>> commodityDataSet = env.fromCollection(commodityList);
        DataSet<Tuple2<String, Integer>> orderDataSet = env.fromCollection(orderList);

        // coGroup操作，两个集合合并为一个数据集
        DataSet<Tuple2<String, String>> coGroupDataSet = commodityDataSet.coGroup(orderDataSet)
                .where((key) -> key.f0)
                .equalTo((key) -> key.f1)
                .with(new CoGroupFunction<Tuple2<Integer, String>, Tuple2<String, Integer>, Tuple2<String, String>>() {
                    @Override
                    public void coGroup(Iterable<Tuple2<Integer, String>> first, Iterable<Tuple2<String, Integer>> second, Collector<Tuple2<String, String>> collector) throws Exception {
                        String names = "";
                        Iterator<Tuple2<String, Integer>> secondIterator = second.iterator();
                        while (secondIterator.hasNext()) {
                            Tuple2<String, Integer> next = secondIterator.next();
                            names = StringUtils.join(names, " ", next.f0);
                        }
                        Iterator<Tuple2<Integer, String>> firstIterator = first.iterator();
                        while (firstIterator.hasNext()) {
                            Tuple2<Integer, String> next = firstIterator.next();
                            collector.collect(new Tuple2<>(next.f1, names));
                        }
                    }
                });
        coGroupDataSet.print("CoGroupFunction");

        env.execute("CoGroupOperator");
    }
}
